#include <string>
#include <sstream>
#include <string.h>
#include "PkgBuf.h"
#include "BufMgr.h"
#include "ModMgr.h"
#include "WebSocketProt.h"
#include "Crypto.h"
#include "Client.h"

#define REQUEST_LINE_END "\r\n"
#define REQUEST_END_FLAG "\r\n\r\n"
#define MAGIC_KEY "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
#define min(x, y) ((x) < (y) ? (x) : (y))

// 打包数据默认不加掩码
bool CWebSocketProt::m_bMask = false;

CWebSocketProt::CWebSocketProt()
{
}

CWebSocketProt::~CWebSocketProt()
{
}

CWebSocketProt *CWebSocketProt::WebSocketProtParse(CBuffer<> *pBuffer, PkgBuf *pbRecv, ProtModule &pm)
{
    bool bRes = false;
    uint32_t nRead = 0;
    char *pEndFlag = nullptr;
    CWebSocketProt *pPort = nullptr;

    do
    {
        bRes = pBuffer->ReadLine(pbRecv->pBuffer + pbRecv->nCurLen, nRead);
        pbRecv->nCurLen += nRead;

        if (bRes)
        {
            pEndFlag = (char *)pbRecv->pBuffer + pbRecv->nCurLen - 4;

            if (0 == strncmp(pEndFlag, REQUEST_END_FLAG, 4))
            {
                pPort = RequestParse(pbRecv, pm);
                break;
            }
        }
    } while (bRes);

    return pPort;
}

int8_t CWebSocketProt::ModuleParse(CBuffer<> *pBuffer, PkgBuf *pbRecv)
{
    return -1;
}

HandleResult CWebSocketProt::DataParse(CBuffer<> *pBuffer, PkgBuf *pbRecv)
{
    int iRecv = 0;
    int iSize = 0;
    uint8_t *pkg = nullptr;

    if (nullptr == pBuffer || nullptr == pbRecv)
        return HandleResult::Error;

    // 还没有取到WS头数据
    if (0 == pbRecv->nPkgLen)
    {
        HandleResult hr = PackageHeaderParse(pBuffer, pbRecv);

        if (HandleResult::Success != hr)
        {
            printf("DataParse-1\n");
            return hr;
        }
    }

    iSize = pBuffer->GetDataLength();

    if ((0 < iSize) && (0 < pbRecv->nPkgLen) && (pbRecv->nCurLen < pbRecv->nPkgLen))
    {
        int iLen = min((pbRecv->nPkgLen - pbRecv->nCurLen), iSize);

        pkg = pbRecv->pBuffer + pbRecv->nCurLen;
        iRecv = pBuffer->Read(pkg, iLen);

        if (-1 == iRecv)
        {
            printf("DataParse-2\n");
            return HandleResult::Error;
        }

        pbRecv->nCurLen += iRecv;

        if (iRecv < iLen)
        {
            printf("DataParse-3\n");
            return HandleResult::Fail;
        }
    }

    if (pbRecv->nCurLen < pbRecv->nPkgLen)
    {
        printf("DataParse-4\n");
        return HandleResult::Fail;
    }

    if ((0 < pbRecv->nPkgLen) && m_wsHeader.bMask)
    {
        pkg = pbRecv->pBuffer;

        for (size_t iLoop = 0; iLoop < pbRecv->nPkgLen; ++iLoop)
            *pkg++ ^= m_wsHeader.szMask[iLoop & 0x03];
    }

    return HandleResult::Success;
}

HandleResult CWebSocketProt::DataHandle(PkgBuf *pbRecv, CClient *pClnt)
{
    HandleResult hr = HandleResult::Fail;

    switch (m_wsHeader.opType)
    {
    case OpCode::TEXT:
    case OpCode::BINARY:
    case OpCode::CONTINUATION:
        break;

    case OpCode::PING:
    {
        printf("OpCode::PING\n");
        int iSent = 0;
        uint16_t nSize = PackageData(pbRecv->pBuffer, OpCode::PONG, nullptr, 0);

        iSent = pClnt->SendPackage(pbRecv->pBuffer, nSize);

        if (iSent >= nSize)
            hr = HandleResult::Success;

        break;
    }

    case OpCode::PONG:
        hr = HandleResult::Success;
        break;

    case OpCode::CLOSE:
        // 返回一个错误，让WorkSvr关闭客户端
        printf("OpCode::CLOSE\n");
        hr = HandleResult::Error;
        break;

    default:
        hr = HandleResult::Error;
        break;
    }

    m_wsHeader = {0};

    return hr;
}

bool CWebSocketProt::ProtResponse(PkgBuf *pbSend, bool bResult)
{
    std::string strHeader;
    auto it = m_mapHeaders.find("Sec-WebSocket-Key");

    if (it == m_mapHeaders.end())
        return false;

    auto sha1 = Crypto::sha1(it->second + MAGIC_KEY);
    strHeader = "HTTP/1.1 101 Switching Protocols\r\n";
    strHeader += "Upgrade: websocket\r\n";
    strHeader += "Connection: Upgrade\r\n";
    strHeader += "Sec-WebSocket-Accept: ";
    strHeader += Crypto::Base64::encode(sha1);
    strHeader += "\r\n\r\n";

    pbSend->nPkgLen = pbSend->nCurLen = strHeader.length();
    memcpy(pbSend->pBuffer, strHeader.c_str(), pbSend->nPkgLen);

    return true;
}

bool CWebSocketProt::ModuleResponse(PkgBuf *pbSend, bool bResult)
{
    return true;
}

uint16_t CWebSocketProt::PackageData(uint8_t *pDest, uint8_t *pSrc, uint16_t nLen)
{
    return PackageData(pDest, OpCode::TEXT, pSrc, nLen);
}

// 服务器发送给客户端，默认不掩码, 最大64K
uint16_t CWebSocketProt::PackageData(uint8_t *pDest, OpCode opType, uint8_t *pSrc, uint64_t nLen)
{
    uint16_t nSize = 0;
    // 掩码 TODO:使用随机数生成
    const uint8_t szMaskKey[4] = {0x12, 0x34, 0x56, 0x78};

    if (nullptr == pDest)
        return 0;

    pDest[0] = 0x80 | (uint8_t)opType;
    nSize++;

    if (nLen < 126)
    {
        // 默认不加掩码
        pDest[nSize++] = (nLen & 0xff) | (m_bMask ? 0x80 : 0);
    }
    else if (nLen < 65536)
    {
        pDest[nSize++] = 126 | (m_bMask ? 0x80 : 0);
        pDest[nSize++] = (nLen >> 8) & 0xff;
        pDest[nSize++] = (nLen >> 0) & 0xff;
    }
    else
    {
        pDest[nSize++] = 127 | (m_bMask ? 0x80 : 0);
        pDest[nSize++] = (nLen >> 56) & 0xff;
        pDest[nSize++] = (nLen >> 48) & 0xff;
        pDest[nSize++] = (nLen >> 40) & 0xff;
        pDest[nSize++] = (nLen >> 32) & 0xff;
        pDest[nSize++] = (nLen >> 24) & 0xff;
        pDest[nSize++] = (nLen >> 16) & 0xff;
        pDest[nSize++] = (nLen >> 8) & 0xff;
        pDest[nSize++] = (nLen >> 0) & 0xff;
    }

    if (m_bMask)
    {
        pDest[nSize++] = szMaskKey[0];
        pDest[nSize++] = szMaskKey[1];
        pDest[nSize++] = szMaskKey[2];
        pDest[nSize++] = szMaskKey[3];
    }

    uint8_t *pData = pDest + nSize;

    if ((nullptr != pSrc) && (0 != nLen))
    {
        memcpy(pData, pSrc, nLen);
        nSize += nLen;

        if (m_bMask)
        {
            for (size_t i = 0; i < nLen; ++i)
                *pData++ ^= szMaskKey[i & 0x3];
        }
    }

    return nSize;
}

char *CWebSocketProt::GetModuleName(char *pBuffer, int &iLen)
{
    char *pEnd = nullptr;
    char *pStart = strchr(pBuffer, ' ');

    if (nullptr != pStart)
    {
        pStart = strchr(pStart, '/');

        if (nullptr != pStart)
        {
            pStart++;

            pEnd = strchr(pStart, ' ');

            if (nullptr != pEnd)
            {
                iLen = pEnd - pStart;
            }
        }
    }

    return pStart;
}

bool CWebSocketProt::GetHeader(char *pStart, char *pEnd, HeadersType &mapHeader)
{
    char *pSep = strchr(pStart, ':');

    if (nullptr != pSep)
    {
        std::string strKey = std::string(pStart, pSep - pStart);
        pSep++;

        while (' ' == *pSep)
            pSep++;

        std::string strVal = std::string(pSep, pEnd - pSep);

        mapHeader[strKey] = strVal;
        // printf("strKey=%s, strVal=%s\n", strKey.c_str(), strVal.c_str());
    }

    return true;
}

CWebSocketProt *CWebSocketProt::RequestParse(PkgBuf *pbRecv, ProtModule &pm)
{
    bool bRes = false;
    char *pStart = (char *)pbRecv->pBuffer;
    char *pEnd = nullptr;
    std::string strModule;
    HeadersType mapHeader;
    CWebSocketProt *pPort = nullptr;

    while (true)
    {
        pEnd = strstr(pStart, REQUEST_LINE_END);

        if (nullptr == pEnd)
            break;

        if (pStart == pEnd)
            break;

        if (strModule.empty())
        {
            int iLen = 0;
            char *pModName = GetModuleName(pStart, iLen);

            if (nullptr != pModName)
                strModule = std::string(pModName, iLen);
        }
        else
        {
            GetHeader(pStart, pEnd, mapHeader);
        }

        pStart = pEnd + strlen(REQUEST_LINE_END);
    }

    auto it = mapHeader.find("Upgrade");

    if (it != mapHeader.end())
    {
        if (0 == strcmp("websocket", it->second.c_str()))
            bRes = true;
    }

    if (bRes && !strModule.empty())
    {
        bRes = false;
        CModMgr *pModMgr = CModMgr::GetInstance();

        if (nullptr == pModMgr)
            return nullptr;

        pm.pt = ProtType::WebSocket;
        pm.iMod = pModMgr->GetModuleId((char *)strModule.c_str());

        if (-1 != pm.iMod)
        {
            pPort = new CWebSocketProt();

            if (nullptr != pPort)
                pPort->m_mapHeaders = mapHeader;

            printf("IsWebSocketProt success\n");
        }
    }

    return pPort;
}

HandleResult CWebSocketProt::PackageHeaderParse(CBuffer<> *pBuffer, PkgBuf *pbRecv)
{
    int iRecv = 0;
    int iSize = pBuffer->GetDataLength();
    uint8_t *pkg = pbRecv->pBuffer;
    HandleResult hr = HandleResult::Fail;

    if (2 > pbRecv->nCurLen)
    {
        if (2 - pbRecv->nCurLen > iSize)
            return HandleResult::Fail;

        iRecv = pBuffer->Read(pkg, 2 - pbRecv->nCurLen);

        if (-1 == iRecv)
            return HandleResult::Error;

        pbRecv->nCurLen += iRecv;
        iSize -= iRecv;

        if (pbRecv->nCurLen >= 2)
        {
            m_wsHeader.bFin = (pkg[0] & 0x80) == 0x80;
            m_wsHeader.opType = (OpCode)(pkg[0] & 0x0f);
            m_wsHeader.bMask = (pkg[1] & 0x80) == 0x80;
            m_wsHeader.nLen = (pkg[1] & 0x7f);
            m_wsHeader.nSize = 2 + (m_wsHeader.nLen == 126 ? 2 : 0) + (m_wsHeader.nLen == 127 ? 8 : 0) + (m_wsHeader.bMask ? 4 : 0);
        }
    }

    if (m_wsHeader.nSize > pbRecv->nCurLen)
    {
        int iLen = 0;
        int iNeedRead = m_wsHeader.nSize - pbRecv->nCurLen;

        if (iNeedRead > iSize)
            return HandleResult::Fail;

        iRecv = pBuffer->Read(pkg + 2, iNeedRead);

        if (-1 == iRecv)
            return HandleResult::Error;

        pbRecv->nCurLen += iRecv;
        iSize -= iRecv;

        if (iRecv < iNeedRead)
            return HandleResult::Error;

        if (m_wsHeader.nLen < 126)
        {
            m_wsHeader.nExtLen = m_wsHeader.nLen;
            iLen = 2;
        }
        else if (m_wsHeader.nLen == 126)
        {
            m_wsHeader.nExtLen = 0;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[2]) << 8;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[3]) << 0;
            iLen = 4;
        }
        else if (m_wsHeader.nLen == 127)
        {
            m_wsHeader.nExtLen = 0;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[2]) << 56;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[3]) << 48;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[4]) << 40;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[5]) << 32;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[6]) << 24;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[7]) << 16;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[8]) << 8;
            m_wsHeader.nExtLen |= ((uint64_t)pkg[9]) << 0;
            iLen = 10;
        }

        if (m_wsHeader.bMask)
        {
            m_wsHeader.szMask[0] = ((uint8_t)pkg[iLen + 0]);
            m_wsHeader.szMask[1] = ((uint8_t)pkg[iLen + 1]);
            m_wsHeader.szMask[2] = ((uint8_t)pkg[iLen + 2]);
            m_wsHeader.szMask[3] = ((uint8_t)pkg[iLen + 3]);
        }
        else
        {
            m_wsHeader.szMask[0] = 0;
            m_wsHeader.szMask[1] = 0;
            m_wsHeader.szMask[2] = 0;
            m_wsHeader.szMask[3] = 0;
        }

        pbRecv->nPkgLen = m_wsHeader.nExtLen;
        // WS头解析完毕，缓冲只存放纯数据
        pbRecv->nCurLen = 0;
        hr = HandleResult::Success;
    }

    return hr;
}